/*
 * covertree_impl.h
 *
 *  Created on: Mar 1, 2011
 *      Author: tqlong
 */

#ifndef COVERTREE_IMPL_H_
#define COVERTREE_IMPL_H_
#include "covertree.h"

#include <algorithm>
#include <cmath>
#include <iostream>
#include <limits>

#define DOUBLE_INFINITY std::numeric_limits<double>::infinity()

template<typename T>
void CoverTreeNode<T>::print()
{
  std::cout << "node lv = " << level_ << " c = " << this->point(0) << "idx =";
  for (int i = 0; i < this->n(); i++)
    std::cout << " " << this->idx_[this->lo_+i];
  std::cout << "\n";
  for (int i = 0; i < this->n_child(); i++)
    child(i)->print();
}

template<typename T>
void CoverTreeNode<T>::prepareRoot()
{
//  std::cerr << "n = " << this->n() << "\n"; std::cerr.flush();
  self_child_ = true;        // root is always a self child
  int rId = ::rand() % this->n();
  this->swapIndex(rId, 0);   // make the first point be center point

  double max = 0;            // compute radius
  for (int i = 1; i < this->n(); i++)
    max = std::max(max, this->distance(0, i));
  assert(max > 0.0);
  level_ = floor(log(max) / log(2) + 1);
//  std::cerr << "lv = " << level_ << " max = " << max
//      << " root = " << this->point(0) << "\n"; std::cerr.flush();
}

template<typename T>
void CoverTreeNode<T>::splitRecursive()
{
  if (this->n() <= 1) return;
  this->split();
  for (int i = 0; i < this->n_child(); i++)
    child(i)->splitRecursive();
}

template<typename T>
void CoverTreeNode<T>::split()
{
//  std::cerr << "split lv = " << level_ << " center = " << this->point(0) << "\n";
  if (level_ < -32) return; // radius too small to split
  // the first point (index 0) is the center
  // radius = 2^level_ cover all points
  // try cover all points by 2^(level_ - 1) using current center
  int i = 1;
  double r = ::pow(2, level_ - 1);
  for (int k = 1; k < this->n(); k++) {
    if (this->distance(0, k) <= r) {
      this->swapIndex(i, k);   // move p_k to beginning of the array
      i++;                     // raise next point index
    }
  }

//  std::cerr << "split lv = " << level_-1 << " c.n = " << i << "\n";

  // check condition to add the self child
  // the self child (if exists) is always the first child
  if (i >= this->n()) {        // cover by 2^(level_ - 1) successful
    if (!self_child_) {        // only add new node if this is not a self child
      CoverTreeNode<T>* c = new CoverTreeNode<T>(*this, 0, i, level_-1, true);
      this->addChild(*c);
//      std::cerr << "!self_child --> add self child\n";
    }
    else {                      // already self child, decrease level
      level_--;
//      std::cerr << "self_child --> make unique\n";
      split();                  // split again at lower level
    }
  }
  else {                       // cover not successful
    CoverTreeNode<T>* c = new CoverTreeNode<T>(*this, 0, i, level_-1, true);
//    std::cerr << "! full cover --> add self child\n";
    this->addChild(*c);
    while (i < this->n()) {
      int r0 = i;
      int rId = ::rand() % (this->n()-i) + i; // choose a random point as center
      this->swapIndex(r0, rId);               // and put it at first position
      i++;
      for (int k = i; k < this->n(); k++)
        if (this->distance(r0, k) <= r) {     // move point to beginning of array
          this->swapIndex(i, k);
          i++;
        }
      CoverTreeNode<T>* c = new CoverTreeNode<T>(*this, r0, i, level_-1, false);
      this->addChild(*c);
//      std::cerr << "split lv = " << level_-1 << " c.n = " << c->n() << "\n";
      assert(i > r0);
    }
  }
}

class Compare
{
public:
  double tol_;
  Compare(double tol = 1e-10) : tol_(tol) {}
  bool equal(double x, double y) { return fabs(x - y) < tol_; }
  bool less(double x, double y) { return x < y-tol_; }
  bool greater(double x, double y) { return x > y+tol_; }
  bool lesseq(double x, double y) { return x <= y + tol_; }
  bool greatereq(double x, double y) { return x >= y-tol_; }
};

//template<typename P, typename CN>
//bool reachLeaf(const P& p, const typename CN::coverset_type &Q_l, int &minIdx)
template<typename T>
bool CoverTreeNode<T>::reachLeaf(const vec_type& p, const coverset_type &Q_l, int &minIdx)
{
  typedef CoverTreeNode<T> CN;
  bool reach_leaf = true;
  minIdx = -1;
  double min = DOUBLE_INFINITY;
  for (unsigned int i = 0; i < Q_l.size(); i++) {
    CN *q_i = Q_l[i];
//    if (q_i->n() > 1 && q_i->level() >= -32) {
    if (q_i->n_child() > 0) {
      reach_leaf = false;
      break;
    }
    double d_pq = arma::norm(p-q_i->center(), 2);
    if (d_pq < min) {
      minIdx = q_i->idx(0);
      min = d_pq;
    }
  }
  return reach_leaf;
}

//template<typename CN>
//bool reachLeafAll(const CN& P,
//               const typename CN::coverset_type &Q_l, std::vector<int>& minIdx)
template<typename T>
bool CoverTreeNode<T>::reachLeafAll(const CoverTreeNode& P, const coverset_type &Q_l, std::vector<int>& minIdx)
{
  typedef CoverTreeNode<T> CN;
  bool reach_leaf = true;
  for (unsigned int i = 0; i < Q_l.size(); i++) {
    CN *q_i = Q_l[i];
//    if (q_i->n() > 1 && q_i->level() >= -32) {
    if (q_i->n_child() > 0) {
      reach_leaf = false;
      break;
    }
  }
  if (!reach_leaf) return false;

  static int count = 0;

  count += P.n();
//  std::cerr << count << " "; std::cerr.flush();
  for (int i = 0; i < P.n(); i++) {                  // for all leaves of P
    int id = P.idx(i);                               // get the id of p[i]
    const typename CN::vec_type& p = P.point(i);     // the point itself
    double min = DOUBLE_INFINITY;                    // find min distance
    int minId = -1;
    for (unsigned int i = 0; i < Q_l.size(); i++) {
      CN *q_i = Q_l[i];
      double d_pq = arma::norm(p-q_i->center(), 2);
      if (d_pq < min) {
        minId = q_i->idx(0);
        min = d_pq;
      }
    }
    minIdx[id] = minId;
  }
  return true;
}

//template<typename P, typename CN>
//double distanceToChild(const P& p, const typename CN::coverset_type &Q_l, double l)
template<typename T>
double CoverTreeNode<T>::distanceToChild(const vec_type& p, const coverset_type &Q_l, double l)
{
  typedef CoverTreeNode<T> CN;
  double d_p_Q = DOUBLE_INFINITY;

  // compute d(p, Q)
  for (unsigned int i = 0; i < Q_l.size(); i++) {
    CN *q_i = Q_l[i];

    double d_p_qi = arma::norm(p-q_i->center(), 2);
    if (d_p_Q > d_p_qi)  // a node is it self-child in implicit representation
      d_p_Q = d_p_qi;

    for (int j = 0; j < q_i->n_child(); j++) {
      CN* q_j = q_i->child(j);
      double d_p_qj = arma::norm(p-q_j->center(), 2);
      if (q_j->level() == l-1 && d_p_Q > d_p_qj)
        d_p_Q = d_p_qj;
    }
  }
  return d_p_Q;
}

//template<typename P, typename CN>
//void nextCoverSet(const P& p, const typename CN::coverset_type &Q_l,
//    double l, double bound,
//    typename CN::coverset_type &Q_lm1)
template<typename T>
void CoverTreeNode<T>::nextCoverSet(const vec_type& p, const coverset_type &Q_l,
                                    double l, double bound, coverset_type &Q_lm1)
{
  typedef CoverTreeNode<T> CN;
  Compare comp;
  for (unsigned int i = 0; i < Q_l.size(); i++) {
    CN *q_i = Q_l[i];

//    std::cerr << "p = " << q_i->idx(0) << "\n";
    if (q_i->n_child() == 0 || q_i->level() != l) {
      double d_p_qi = arma::norm(p-q_i->center(), 2);
//      std::cerr << "check c = " << q_i->idx(0) << " = " << d_p_qi << "\n";
      if (comp.greatereq(bound, d_p_qi))  // a node is it self-child in implicit representation
        Q_lm1.push_back(q_i);
    }
    else {
      for (int j = 0; j < q_i->n_child(); j++) {
        CN* q_j = q_i->child(j);
        double d_p_qj = arma::norm(p-q_j->center(), 2);
//        std::cerr << "check c = " << q_j->idx(0) << " = " << d_p_qj << "\n";
        if (comp.greatereq(bound, d_p_qj))
          Q_lm1.push_back(q_j);
      }
    }
  }
}

//template<typename P, typename CN>
//int nearest(const P& p,
//            const typename CN::coverset_type &Q_l, double l)
template<typename T>
int CoverTreeNode<T>::nearest(const vec_type& p, const coverset_type &Q_l, double l)
{
  typedef CoverTreeNode<T> CN;
  int min_idx;
  bool reach_leaf = reachLeaf (p, Q_l, min_idx);
  //  std::cerr << "lv = " << l << " Ql.size = " << Q_l.size() << " "
  //      << (!reach_leaf ? "not reach leaf":"reach leaf") << "\n";
  //  std::cerr << "nodes =";
  //  for (unsigned int i = 0; i < Q_l.size(); i++)
  //    std::cerr << " " << Q_l[i]->idx(0);
  //  std::cerr << "\n";
  if (reach_leaf)
    return min_idx;

  typename CN::coverset_type Q_lm1;

  // compute Q(l-1) = { q : d(p, q) <= 2^l + d(p, Q) };
  double bound = distanceToChild (p, Q_l, l) + ::pow(2, l);
//  std::cerr << "bound = " << bound << "\n";
  nextCoverSet (p, Q_l, l, bound, Q_lm1);
  return nearest (p, Q_lm1, l-1);
}

//template<typename P, typename CN>
//int nearest(const P& p,
//            CN& root)
template<typename T>
int CoverTreeNode<T>::nearest(const vec_type& p, CoverTreeNode& root)
{
  typedef CoverTreeNode<T> CN;
  typename CN::coverset_type Ql;
  Ql.push_back(&root);
  return nearest (p, Ql, root.level());
}

//template<typename CN>
//void allNearest(const CN& P, const typename CN::coverset_type &Q_l,
//               double lP, double lQ, std::vector<int>& minIdx)
template<typename T>
void CoverTreeNode<T>::allNearest(const CoverTreeNode& P, const coverset_type &Q_l,
                         double lP, double lQ, std::vector<int>& minIdx)
{
  typedef CoverTreeNode<T> CN;
  typedef typename CN::vec_type vec_type;
  bool reach_leaf = reachLeafAll (P, Q_l, minIdx);
//  std::cerr << "lv = (" << lP << ", " << lQ << ") Ql.size = " << Q_l.size() << " "
//      << (!reach_leaf ? "not reach leaf":"reach leaf") << "\n";
//  std::cerr << "nodes =";
//  for (unsigned int i = 0; i < Q_l.size(); i++)
//    std::cerr << " " << Q_l[i]->idx(0);
//  std::cerr << "\n";
//  std::cerr.flush();
  if (reach_leaf)
    return;

  if (lP < lQ) {
    typename CN::coverset_type Q_lm1;
    // compute Q(l-1) = { q : d(p, q) <= 2^lQ + 2^(lP+2) + d(p, Q) };
    double bound = distanceToChild(P.center(), Q_l, lQ)
        + ::pow(2, lQ) + ::pow(2, lP+2);
    //  std::cerr << "bound = " << bound << "\n";
    nextCoverSet (P.center(), Q_l, lQ, bound, Q_lm1);
    allNearest (P, Q_lm1, lP, lQ-1, minIdx);
  }
  else {
    if (P.n_child() == 0 || P.level() != lP)
      allNearest (P, Q_l, lP-1, lQ, minIdx);
    else {
      for (int i = 0; i < P.n_child(); i++) {
        CN* p_i = P.child(i);
        allNearest (*p_i, Q_l, lP-1, lQ, minIdx);
      }
    }
  }
}

//template<typename CN>
//void allNearest(CN& P, CN& Q, std::vector<int>& minIdx)
template<typename T>
void CoverTreeNode<T>::allNearest(CoverTreeNode& P, CoverTreeNode& Q, std::vector<int>& minIdx)
{
  typedef CoverTreeNode<T> CN;
  typename CN::coverset_type Ql;
  Ql.push_back(&Q);
  minIdx = std::vector<int>(P.n(), -1);
  allNearest(P, Ql, P.level(), Q.level(), minIdx);
}
#endif /* COVERTREE_IMPL_H_ */
